from mcp.server.fastmcp import FastMCP
from utils import logger
from utils.question_match import common_question, statis_question
from utils.config_load import load_prompt, load_database, load_host_port
from utils.llm import llm_api
from utils.sql_process import generate_sql, fix_sql
from schema import db_connect
import sqlparse
import argparse


mcp = FastMCP("自然语言与SQL转换工具")
dialect, database = load_database()


@mcp.resource(dialect + '://' + database)
async def read_resource() -> str:
    """获取数据库中所有表格信息"""
    return db_connect.get_schema()


@mcp.resource(dialect + "://{table_name}")
async def read_resource(table_name: str) -> str:
    """获取数据库中，指定表格信息"""
    return db_connect.get_schema(table_name)


@mcp.tool()
def get_database_info() -> str:
    """
    获取数据库的全部表结构及案例数据
    :return: 数据库全部表格的结构信息
    """
    return db_connect.get_schema()


@mcp.tool()
def get_table_info(table_name: str) -> str:
    """
    查询数据库中的指定表格
    :param table_name: 待查询表名
    :return: 指定表格的结构信息
    """
    return db_connect.get_schema(table_name)


@mcp.tool()
def text_to_sql(question: str) -> str:
    """
    根据用户输入的自然语言，生成对应的 SQL 查询
    :param question: 用户提问
    :return: 可执行 SQL
    """
    # Step 1. 匹配常见问题
    com_ques_match_res = common_question(question)
    cq_match_res = str(list(com_ques_match_res.values())[0])
    if list(com_ques_match_res.keys())[0] == 'success':
        return cq_match_res
    # Step 2. 匹配统计问题
    stat_ques_match_res = statis_question(question)
    if list(stat_ques_match_res.keys())[0] == 'success':
        return str(list(stat_ques_match_res.values())[0])
    # Step 3. 调用 LLM API 生成 SQL
    sql = generate_sql(question=question, case_qa=cq_match_res)
    if sql == "None":
        return "目前暂无相关数据"
    else:
        # Step 4. 检验 SQL 的正确性
        sql, _ = fix_sql(sql, question, cq_match_res)
        if sql:
            return sql
        else:
            return "目前暂无相关数据"


@mcp.tool()
def sql_execute_result(question: str, sql: str) -> str:
    """
    SQL 执行结果
    :param question: 用户提问
    :param sql: 可执行 SQL
    :return: SQL 执行结果
    """
    sql_result = "目前暂无相关数据"
    if sql in ("None", "无法生成 SQL", "目前暂无相关数据"):
        logger.info(f"【execute sql】{sql} 【result】 {sql_result}")
        return sql_result
    else:
        # 仅执行 SELECT 语句
        if sqlparse.parse(sql)[0].get_type() == 'SELECT':
            execute_result = db_connect.execute_sql(sql)
            if list(execute_result.keys())[0] == 'fail':
                logger.info(f"【execute sql】{sql} 【result】 {sql_result}")
                return sql_result
            else:
                sql_result = str(list(execute_result.values())[0])
                logger.info(f"【execute sql】{sql} 【result】 {sql_result}")
                # 将 SQL 执行结果转换为 markdown 格式
                sql_md = llm_api(load_prompt(user_input=question, task='sql2md', sql=sql, sql_result=sql_result))
                return sql_md
        else:
            logger.info(f"【execute sql】{sql} 【result】 禁止执行非查询 SQL")
            return sql_result


@mcp.tool()
def text_query_database(question: str) -> str:
    """
    根据用户输入的自然语言，查询数据库中的对应结果
    :param question: 用户提问
    :return: 数据库查询结果
    """
    sql = text_to_sql(question)
    if sql == "目前暂无相关数据":
        logger.info(f"fail 【generate sql】 无法生成 SQL")
        return "目前暂无相关数据"
    else:
        sql_result = sql_execute_result(question, sql)
        if sql_result == "目前暂无相关数据":
            logger.info(f"【execute sql】{sql} 【result】 {sql_result}")
        else:
            logger.info(f"【execute sql】{sql} 【result】 {sql_result}")
        return sql_result


@mcp.tool()
def expansion_questions(question: str) -> dict:
    """
    根据用户提问和数据库信息，扩展成多个不同维度、且与之相关的拓展问题，并返回对应的数据库查询结果
    :param question: 用户提问
    :return: 拓展问题与对应的 SQL查询结果
    """
    com_ques_match_res = common_question(question)
    case_qa = str(list(com_ques_match_res.values())[0])
    questions = llm_api(load_prompt(user_input=question, task='expansion_question', case_qa=case_qa))
    try:
        questions = questions.split(',')
        logger.info(f"success 【expansion questions】 {str(questions)}")
        question_list = []
        answer_list = []
        for question in questions:
            sql = text_to_sql(question)
            if sql == "目前暂无相关数据":
                logger.info(f"fail 【generate sql】 无法生成 SQL")
            else:
                sql_result = sql_execute_result(question, sql)
                if sql_result == "目前暂无相关数据":
                    logger.info(f"【execute sql】{sql} 【result】 {sql_result}")
                else:
                    logger.info(f"【execute sql】{sql} 【result】 {sql_result}")
                    question_list.append(question)
                    answer_list.append(sql_result)
        return dict(zip(question_list, answer_list))
    except Exception as e:
        logger.error(f"fail 【expansion questions】 {str(e)}")
        return {"无法生成相关问题": "目前暂无相关数据"}


def main():
    parser = argparse.ArgumentParser(description="Run MCP server.")
    parser.add_argument('transport', nargs='?', default='stdio', choices=['stdio', 'sse'],
                        help='Transport type (stdio or sse)')
    args = parser.parse_args()
    mcp.run(transport=args.transport)
    # 开始记录日志
    logger.info("text2sql mcp server start!!!")


if __name__ == "__main__":
    print(text_to_sql("your question"))
    print(sql_execute_result(question="your question", sql="corresponding SQL"))
    print(expansion_questions("your question"))
    main()
